from alphasql.algorithm.mcts.mcts_node import MCTSNode, MCTSNodeType
from alphasql.algorithm.mcts.reward import RewardModel
from alphasql.runner.task import Task
from alphasql.llm_call.openai_llm import call_openai
import random
import copy
from typing import List, Dict, Any
from pathlib import Path
import pickle
from tqdm import tqdm

class EvolutionarySolver:
    def __init__(self,
                 db_root_dir: str,
                 task: Task,
                 population_size: int,
                 generations: int,
                 mutation_rate: float,
                 save_root_dir: str,
                 llm_kwargs: Dict[str, Any],
                 reward_model: RewardModel):
        self.llm_kwargs = llm_kwargs
        self.reward_model = reward_model
        self.task = task
        self.db_root_dir = db_root_dir
        self.population_size = population_size
        self.generations = generations
        self.mutation_rate = mutation_rate
        self.save_root_dir = save_root_dir
        self.visited_states = set()  # For pruning

    def initialize_population(self) -> List[MCTSNode]:
        """Initialize population with LLM-generated SQL queries"""
        population = []
        for _ in range(self.population_size):
            root = MCTSNode(
                node_type=MCTSNodeType.ROOT,
                original_question=self.task.question,
                selected_schema_dict=self.task.table_schema_dict or {},
                rephrased_question=self.task.question,  # Simplified
                llm_kwargs=self.llm_kwargs
            )

            # Generate SQL using LLM
            sql_prompt = f"Generate a SQL query for: {self.task.question}\nSchema: {self.task.table_schema_dict or {}}"
            sql_response = call_openai(sql_prompt, **self.llm_kwargs)
            sql_query = sql_response[0] if isinstance(sql_response, list) else sql_response

            # Create SQL generation node
            sql_node = MCTSNode(
                node_type=MCTSNodeType.SQL_GENERATION,
                parent_node=root,
                original_question=self.task.question,
                selected_schema_dict=self.task.table_schema_dict or {},
                sql_query=sql_query,
                llm_kwargs=self.llm_kwargs
            )
            root.children.append(sql_node)
            population.append(sql_node)

        return population

    def evaluate_fitness(self, individual: MCTSNode) -> float:
        """Evaluate fitness using reward model"""
        if individual.node_type == MCTSNodeType.END:
            return self.reward_model.get_reward(individual)
        return 0.0

    def mutate(self, individual: MCTSNode) -> MCTSNode:
        """Mutate an individual by generating a variant SQL"""
        if random.random() < self.mutation_rate:
            mutation_prompt = f"Modify this SQL query to make it better or fix potential issues: {individual.sql_query}\nQuestion: {individual.original_question}"
            mutated_sql = call_openai(mutation_prompt, **self.llm_kwargs)
            mutated_sql = mutated_sql[0] if isinstance(mutated_sql, list) else mutated_sql

            # Create new node with mutated SQL
            mutated_node = copy.deepcopy(individual)
            mutated_node.sql_query = mutated_sql
            return mutated_node
        return individual

    def select_parents(self, population: List[MCTSNode], fitness_scores: List[float]) -> List[MCTSNode]:
        """Select parents using tournament selection"""
        parents = []
        for _ in range(2):  # Select 2 parents
            # Tournament of 3
            candidates = random.sample(list(zip(population, fitness_scores)), 3)
            winner = max(candidates, key=lambda x: x[1])[0]
            parents.append(winner)
        return parents

    def crossover(self, parent1: MCTSNode, parent2: MCTSNode) -> MCTSNode:
        """Simple crossover by combining SQL parts"""
        # For simplicity, just pick one parent's SQL or create a hybrid
        if random.random() < 0.5:
            return copy.deepcopy(parent1)
        else:
            return copy.deepcopy(parent2)

    def evolve(self) -> MCTSNode:
        """Main evolutionary algorithm"""
        print(f"\n🧬 初始化种群 (大小: {self.population_size})...")
        population = self.initialize_population()

        print(f"\n🧬 进化算法进度 (任务 {self.task.question_id}):")
        for generation in tqdm(range(self.generations), desc="进化代数", unit="gen", ncols=100):
            # Evaluate fitness
            fitness_scores = [self.evaluate_fitness(ind) for ind in population]

            # Select best individuals
            best_indices = sorted(range(len(fitness_scores)), key=lambda i: fitness_scores[i], reverse=True)
            elite_size = max(1, self.population_size // 4)
            elites = [population[i] for i in best_indices[:elite_size]]

            # Create new population
            new_population = elites.copy()

            while len(new_population) < self.population_size:
                # Select parents
                parents = self.select_parents(population, fitness_scores)

                # Crossover
                offspring = self.crossover(parents[0], parents[1])

                # Mutate
                offspring = self.mutate(offspring)

                new_population.append(offspring)

            population = new_population

            # Print best fitness
            best_fitness = max(fitness_scores)
            tqdm.write(f"  第 {generation + 1} 代最佳适应度: {best_fitness:.4f}")

        # Return best individual
        fitness_scores = [self.evaluate_fitness(ind) for ind in population]
        best_index = fitness_scores.index(max(fitness_scores))
        best_individual = population[best_index]

        # Create END node
        end_node = MCTSNode(
            node_type=MCTSNodeType.END,
            parent_node=best_individual,
            original_question=best_individual.original_question,
            selected_schema_dict=best_individual.selected_schema_dict or {},
            final_sql_query=best_individual.sql_query,
            llm_kwargs=self.llm_kwargs
        )
        best_individual.children.append(end_node)

        return end_node

    def solve(self) -> MCTSNode:
        """Solve the task using evolutionary algorithm"""
        best_node = self.evolve()

        # Save result similar to MCTS
        save_path = Path(self.save_root_dir) / f"{self.task.question_id}.pkl"
        print(f"Evolutionary solver done for task {self.task.question_id}, saving to {save_path}")

        # Create a list with the best reasoning path (similar to MCTS format)
        reasoning_path = [best_node]  # Just the end node for evolutionary
        with open(save_path, "wb") as f:
            pickle.dump(reasoning_path, f)

        return best_node
